
from psychsim.pwl import *
from psychsim.reward import *
from psychsim.action import *
from psychsim.world import World,stateKey,actionKey,binaryKey,modelKey
from psychsim.agent import Agent

''' 
This example is used to test the update of beliefs on world states
according to the outcome of the last action
'''

class AgentOutcomeTestCase(unittest.TestCase):

    def setUp(self):
        #init world and agent, set maxRounds and turn order
        self.world = World()
        self.claudius = Agent('Claudius')
        self.world.addAgent(self.claudius)

        self.maxRounds=4
        self.world.setOrder(['Claudius'])

    def addStates(self):      
        #add states for round and money
        self.world.defineState(None,'round',int)
        self.world.setState(None,'round',0)

        self.world.defineState(self.claudius.name,'money',int)
        self.claudius.setState('money',10)        
                                    
    def addTermination(self):
        #add termination if at max round  
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None,'round'),self.maxRounds),
                                   True: True, False: False}))
        
    def addActions(self):
        #add actions and dynamics     
        action =  self.claudius.addAction({'verb': 'spend money'})
        tree_round = makeTree(incrementMatrix(stateKey(None,'round'),1))
        self.world.setDynamics('round',action,tree_round)
        tree = makeTree(incrementMatrix(stateKey('Claudius','money'),-1))
        self.world.setDynamics(stateKey('Claudius', 'money'),action,tree)
        
        action =  self.claudius.addAction({'verb': 'do nothing'})
        tree_round = makeTree(incrementMatrix(stateKey(None,'round'),1))
        self.world.setDynamics('round',action,tree_round)
         
                    
    def specifyModel(self):
        #make sure to update model
        self.world.setModel('Claudius', True)
        
        key = stateKey('Claudius','money')
        self.claudius.setBelief(key,Distribution({8:0.5, 5:0.5}),True)

    def addRewards(self):
        #spend money
        self.claudius.setReward(minimizeFeature(stateKey('Claudius','money')), 1., True)

    def testSetup(self):
        self.addStates()
        self.addTermination()
        self.addActions()
        self.specifyModel() 
        self.addRewards()                                                             
            
    def testRun(self):
        self.testSetup()
        self.world.printState(beliefs=True)
        self.world.printModels()
        for t in range(self.maxRounds + 1):
            self.world.explain(self.world.step(),level=3)
            self.world.printState(beliefs=True)
            self.world.printModels()
            if self.world.terminated():
                break 
            

if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(AgentOutcomeTestCase):
       if method.startswith("test"):
          suite.addTest(AgentOutcomeTestCase(method))
    unittest.TextTestRunner().run(suite)